#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
北证50指数成分股数据抓取脚本
基于测试结果，采用智能获取策略：
1. 按季度获取数据
2. 检测变化点，获取详细变化数据
3. 支持增量更新
"""

import sqlite3
import requests
import json
import pandas as pd
from datetime import datetime, timedelta
import logging
import time
from typing import List, Dict, Optional
import os

# 配置
TOKEN = "4bb4c74b-3168-4d7a-ae2e-d679a32c7921"
API_BASE_URL = "https://open.lixinger.com/api/cn/index/constituents"
DB_NAME = "bse_data.db"
REQUEST_DELAY = 0.1

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('bse50_constituents_fetch.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


class BSE50ConstituentsFetcher:
    """北证50指数成分股数据抓取器"""

    def __init__(self, db_path: str = DB_NAME):
        self.db_path = db_path
        self.session = requests.Session()
        self.session.headers.update({
            'Content-Type': 'application/json',
            'User-Agent': 'BSE50-Constituents-Fetcher/1.0'
        })
        self.index_code = "899050"  # 北证50指数代码

        # 计算5年日期范围
        self.end_date = datetime.now().strftime('%Y-%m-%d')
        self.start_date = (datetime.now() - timedelta(days=365*5)).strftime('%Y-%m-%d')

        logger.info(f"数据抓取时间范围: {self.start_date} 至 {self.end_date}")
        self.init_database()

    def init_database(self):
        """初始化数据库表结构"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            # 北证50成分股变化表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS bse50_constituents_changes (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    date DATE NOT NULL,
                    constituents_count INTEGER NOT NULL,
                    constituents TEXT NOT NULL,  -- JSON格式存储成分股列表
                    added_stocks TEXT,           -- JSON格式存储新增股票列表
                    removed_stocks TEXT,         -- JSON格式存储移除股票列表
                    change_type TEXT NOT NULL,   -- 'initial', 'quarterly', 'irregular'
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(date)
                )
            ''')

            # 北证50成分股快照表（每日或关键日期的完整数据）
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS bse50_constituents_snapshot (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    date DATE NOT NULL,
                    stock_code TEXT NOT NULL,
                    stock_name TEXT,
                    weight REAL,  -- 权重（如果API提供）
                    is_active BOOLEAN DEFAULT 1,  -- 是否为当前成分股
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(date, stock_code)
                )
            ''')

            # 创建索引
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_constituents_changes_date ON bse50_constituents_changes(date)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_constituents_snapshot_date ON bse50_constituents_snapshot(date)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_constituents_snapshot_stock ON bse50_constituents_snapshot(stock_code)')

            conn.commit()
            conn.close()
            logger.info("数据库初始化完成")

        except Exception as e:
            logger.error(f"数据库初始化失败: {e}")
            raise

    def _make_request(self, data: Dict, max_retries: int = 3) -> Optional[Dict]:
        """发送API请求"""
        for attempt in range(max_retries):
            try:
                response = self.session.post(API_BASE_URL, json=data, timeout=30)
                response.raise_for_status()

                result = response.json()
                if result.get('code') == 1:
                    return result.get('data')
                else:
                    logger.error(f"API返回错误: {result.get('message', '未知错误')}")
                    return None

            except requests.exceptions.RequestException as e:
                logger.warning(f"请求失败 (尝试 {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    time.sleep(2 ** attempt)
                else:
                    logger.error(f"请求最终失败")
                    return None

            time.sleep(REQUEST_DELAY)

    def get_constituents_on_date(self, date: str) -> List[str]:
        """获取指定日期的北证50成分股"""
        data = {
            "token": TOKEN,
            "date": date,
            "stockCodes": [self.index_code]
        }

        result_data = self._make_request(data)
        if result_data and len(result_data) > 0:
            constituents = result_data[0].get('constituents', [])
            return [item['stockCode'] for item in constituents]

        return []

    def get_quarter_end_dates(self) -> List[str]:
        """获取过去5年的季度末日期列表"""
        dates = []
        current_date = datetime.now()

        # 生成过去20个季度的季度末日期
        for quarter in range(20):
            # 计算季度
            year_offset = quarter // 4
            quarter_num = 3 - (quarter % 4)  # 从Q4开始往前推

            year = current_date.year - year_offset

            # 计算季度末日期
            if quarter_num == 0:  # Q1
                quarter_end = datetime(year, 3, 31)
            elif quarter_num == 1:  # Q2
                quarter_end = datetime(year, 6, 30)
            elif quarter_num == 2:  # Q3
                quarter_end = datetime(year, 9, 30)
            else:  # Q4
                quarter_end = datetime(year, 12, 31)

            # 确保不超过当前日期
            if quarter_end <= current_date:
                dates.append(quarter_end.strftime('%Y-%m-%d'))

        return sorted(dates)

    def get_existing_dates(self) -> set:
        """获取数据库中已存在的日期"""
        conn = sqlite3.connect(self.db_path)
        try:
            cursor = conn.cursor()
            cursor.execute("SELECT date FROM bse50_constituents_changes")
            return {row[0] for row in cursor.fetchall()}
        finally:
            conn.close()

    def detect_changes(self, prev_constituents: List[str], curr_constituents: List[str]) -> Dict:
        """检测成分股变化"""
        prev_set = set(prev_constituents) if prev_constituents else set()
        curr_set = set(curr_constituents)

        added = list(curr_set - prev_set)
        removed = list(prev_set - curr_set)

        return {
            'added': sorted(added),
            'removed': sorted(removed),
            'total_changes': len(added) + len(removed)
        }

    def save_constituents_data(self, date: str, constituents: List[str],
                             prev_constituents: List[str] = None,
                             change_type: str = 'regular'):
        """保存成分股数据"""
        conn = sqlite3.connect(self.db_path)
        try:
            cursor = conn.cursor()

            # 检测变化
            changes = self.detect_changes(prev_constituents, constituents)

            # 保存到变化表
            cursor.execute('''
                INSERT OR REPLACE INTO bse50_constituents_changes
                (date, constituents_count, constituents, added_stocks, removed_stocks, change_type)
                VALUES (?, ?, ?, ?, ?, ?)
            ''', (
                date,
                len(constituents),
                json.dumps(constituents, ensure_ascii=False),
                json.dumps(changes['added'], ensure_ascii=False) if changes['added'] else None,
                json.dumps(changes['removed'], ensure_ascii=False) if changes['removed'] else None,
                change_type
            ))

            # 保存到快照表
            # 先删除该日期的旧数据
            cursor.execute("DELETE FROM bse50_constituents_snapshot WHERE date = ?", (date,))

            # 插入新的快照数据
            for stock_code in constituents:
                cursor.execute('''
                    INSERT INTO bse50_constituents_snapshot
                    (date, stock_code, is_active)
                    VALUES (?, ?, ?)
                ''', (date, stock_code, 1))

            conn.commit()

            # 记录日志
            if changes['total_changes'] > 0:
                logger.info(f"保存 {date} 数据: {len(constituents)} 只成分股, 变化: +{len(changes['added'])} -{len(changes['removed'])}")
            else:
                logger.debug(f"保存 {date} 数据: {len(constituents)} 只成分股, 无变化")

        except Exception as e:
            logger.error(f"保存数据失败 {date}: {e}")
            conn.rollback()
        finally:
            conn.close()

    def fetch_quarterly_data(self):
        """获取季度数据"""
        logger.info("开始获取北证50季度成分股数据...")

        quarter_dates = self.get_quarter_end_dates()
        existing_dates = self.get_existing_dates()

        # 过滤已存在的日期
        dates_to_fetch = [date for date in quarter_dates if date not in existing_dates]

        if not dates_to_fetch:
            logger.info("所有季度数据已存在")
            return

        logger.info(f"需要获取 {len(dates_to_fetch)} 个季度的数据")

        prev_constituents = None
        for i, date in enumerate(dates_to_fetch):
            logger.info(f"获取季度数据 {i+1}/{len(dates_to_fetch)}: {date}")

            constituents = self.get_constituents_on_date(date)
            if constituents:
                # 判断变化类型
                change_type = 'quarterly'
                if not prev_constituents:
                    change_type = 'initial'
                elif self.detect_changes(prev_constituents, constituents)['total_changes'] > 10:
                    change_type = 'major_quarterly'

                self.save_constituents_data(date, constituents, prev_constituents, change_type)
                prev_constituents = constituents
            else:
                logger.warning(f"未能获取 {date} 的数据")

            # 请求间隔
            time.sleep(REQUEST_DELAY)

    def fetch_irregular_changes(self):
        """获取非规则变化数据（基于已知的变化日期）"""
        logger.info("检查非规则成分股变化...")

        # 基于测试结果，我们知道2025-10-08有大规模变化
        # 可以扩展这个列表基于历史数据
        known_change_dates = [
            "2025-10-08",  # 从测试中发现的变化日期
            # 可以添加其他已知的变化日期
        ]

        existing_dates = self.get_existing_dates()
        dates_to_check = [date for date in known_change_dates if date not in existing_dates]

        if not dates_to_check:
            logger.info("没有需要检查的非规则变化日期")
            return

        for date in dates_to_check:
            logger.info(f"检查非规则变化日期: {date}")

            # 获取变化前后的数据
            change_date = datetime.strptime(date, '%Y-%m-%d')
            before_date = (change_date - timedelta(days=1)).strftime('%Y-%m-%d')
            after_date = (change_date + timedelta(days=1)).strftime('%Y-%m-%d')

            # 获取变化前的数据
            before_constituents = self.get_constituents_on_date(before_date)
            if before_constituents:
                self.save_constituents_data(before_date, before_constituents, change_type='irregular_before')

            # 获取变化当天的数据
            current_constituents = self.get_constituents_on_date(date)
            if current_constituents:
                self.save_constituents_data(date, current_constituents, before_constituents, 'irregular')

            # 获取变化后的数据
            after_constituents = self.get_constituents_on_date(after_date)
            if after_constituents:
                self.save_constituents_data(after_date, after_constituents, current_constituents, 'irregular_after')

    def get_data_summary(self) -> Dict:
        """获取数据摘要"""
        conn = sqlite3.connect(self.db_path)
        try:
            cursor = conn.cursor()

            # 基本统计
            cursor.execute("SELECT COUNT(*) FROM bse50_constituents_changes")
            total_changes = cursor.fetchone()[0]

            cursor.execute("SELECT MIN(date), MAX(date) FROM bse50_constituents_changes")
            date_range = cursor.fetchone()

            cursor.execute("SELECT COUNT(DISTINCT stock_code) FROM bse50_constituents_snapshot")
            unique_stocks = cursor.fetchone()[0]

            # 变化统计
            cursor.execute("""
                SELECT change_type, COUNT(*)
                FROM bse50_constituents_changes
                GROUP BY change_type
            """)
            change_types = dict(cursor.fetchall())

            # 最近的变化
            cursor.execute("""
                SELECT date, added_stocks, removed_stocks
                FROM bse50_constituents_changes
                WHERE added_stocks IS NOT NULL OR removed_stocks IS NOT NULL
                ORDER BY date DESC
                LIMIT 5
            """)
            recent_changes = cursor.fetchall()

            return {
                'total_change_records': total_changes,
                'date_range': date_range,
                'unique_constituents': unique_stocks,
                'change_types': change_types,
                'recent_changes': recent_changes
            }

        finally:
            conn.close()

    def run(self):
        """执行完整的数据抓取流程"""
        logger.info("开始北证50成分股数据抓取...")
        start_time = datetime.now()

        try:
            # 1. 获取季度数据
            self.fetch_quarterly_data()

            # 2. 检查非规则变化
            self.fetch_irregular_changes()

            # 3. 输出统计信息
            end_time = datetime.now()
            duration = end_time - start_time
            summary = self.get_data_summary()

            logger.info("=" * 50)
            logger.info("北证50成分股数据抓取完成!")
            logger.info(f"耗时: {duration}")
            logger.info(f"变化记录数: {summary['total_change_records']}")
            logger.info(f"数据时间范围: {summary['date_range'][0]} 至 {summary['date_range'][1]}")
            logger.info(f"历史成分股总数: {summary['unique_constituents']}")
            logger.info(f"变化类型分布: {summary['change_types']}")
            logger.info(f"数据库文件: {os.path.abspath(self.db_path)}")
            logger.info("=" * 50)

        except Exception as e:
            logger.error(f"数据抓取过程中发生错误: {e}")
            raise


def main():
    """主函数"""
    fetcher = BSE50ConstituentsFetcher()
    fetcher.run()


if __name__ == "__main__":
    main()